package com.malty.gateway.config;

import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.function.BiFunction;

import static com.malty.gateway.config.constant.GateWayConstant.GATEWAY_CONTEXT;

/**
 * @author malty
 */
@Slf4j
public class DecoratedResponse extends ServerHttpResponseDecorator {
    private ServerWebExchange exchange;
    private BiFunction<Long, Long, Mono<Boolean>> consumer;

    public DecoratedResponse(ServerWebExchange exchange, BiFunction<Long, Long, Mono<Boolean>> consumer) {
        super(exchange.getResponse());
        this.exchange = exchange;
        this.consumer = consumer;
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        return DataBufferUtils.join(Flux.from(body))
                .flatMap(dataBuffer -> {
                    byte[] content = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(content);
                    DataBufferUtils.release(dataBuffer);

                    logResponseData(content, getDelegate(), exchange);
                    GatewayContext context = exchange.getAttribute(GATEWAY_CONTEXT);
                    Mono<Void> statisticsMono = Mono
                            .defer(() -> consumer.apply(context.getApiId(), context.getUserId()))
                            .then();

                    return super.writeWith(Mono.just(bufferFactory().wrap(content)))
                            .then(statisticsMono);
                });
    }

    /**
     * 记录响应数据日志。
     *
     * @param content          响应体字节数组
     * @param originalResponse 原始的 ServerHttpResponse
     * @param exchange         ServerWebExchange 对象
     */
    private void logResponseData(byte[] content, ServerHttpResponse originalResponse, ServerWebExchange exchange) {
        String responseData = new String(content, StandardCharsets.UTF_8);
        log.info("API Response [Path: {}, Status: {}]: {}",
                exchange.getRequest().getPath(), originalResponse.getStatusCode(), responseData);
    }
}
